#include "LightShader.h"
#include "Device.h"
#include "Light.h"
#include "Camera.h"
#include "Buffer.h"
#include "Utility.h"
#include "Texture.h"
#include "NeroCore\NMemoryManager.h"

#include <d3dx11async.h>

using namespace Graphics;
RTTI_DEFINITIONS(LightShader)

LightShader::LightShader(void)
	:Shader("LightShader",""),
	 mLightBuffer(NULL),
	 mCameraBuffer(NULL)
{
}

LightShader::~LightShader(void)
{
	Deinitialize();
}

bool LightShader::Initialize()
{
	Graphics::D3DDevice* device = GetPlatform()->GetDevice()->As<Graphics::D3DDevice>();
	
	bool result =  InitializeShader(device->GetWindowsSettings().WindowHandle, 
									"..\\..\\assets\\shaders\\light.vs",
									"..\\..\\assets\\shaders\\light.ps");
	return result;
}

void LightShader::Deinitialize()
{
	DeinitializeShader();
}

bool LightShader::Render(int indexCount, D3DXMATRIX worldMatrix, D3DXMATRIX viewMatrix, D3DXMATRIX projectionMatrix, Camera* camera,  Light* light, Texture* texture)
{
	Graphics::D3DDevice* nDevice = GetPlatform()->GetDevice()->As<Graphics::D3DDevice>();
	
	bool result = SetShaderParameters(worldMatrix, viewMatrix, projectionMatrix, camera, light, texture);

	result = RenderShader(nDevice, indexCount);

	return result;
}

bool LightShader::InitializeShader(HWND hwnd, const char* vsFilename, const char* psFilename)
{
	Graphics::D3DDevice* nDevice = GetPlatform()->GetDevice()->As<Graphics::D3DDevice>();
	if(Shader::InitializeShader(nDevice, hwnd, vsFilename, psFilename))
	{
		// Now setup the layout of the data that goes into the shader.
		// This setup needs to match the VertexType stucture in the ModelClass and in the shader.
		D3D11_INPUT_ELEMENT_DESC polygonLayout[4];
		polygonLayout[0].SemanticName = "POSITION";
		polygonLayout[0].SemanticIndex = 0;
		polygonLayout[0].Format = DXGI_FORMAT_R32G32B32_FLOAT;
		polygonLayout[0].InputSlot = 0;
		polygonLayout[0].AlignedByteOffset = 0;
		polygonLayout[0].InputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;
		polygonLayout[0].InstanceDataStepRate = 0;

		polygonLayout[1].SemanticName = "COLOR";
		polygonLayout[1].SemanticIndex = 0;
		polygonLayout[1].Format = DXGI_FORMAT_R32G32B32A32_FLOAT;
		polygonLayout[1].InputSlot = 0;
		polygonLayout[1].AlignedByteOffset = D3D11_APPEND_ALIGNED_ELEMENT;
		polygonLayout[1].InputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;
		polygonLayout[1].InstanceDataStepRate = 0;

		polygonLayout[2].SemanticName = "TEXCOORD";
		polygonLayout[2].SemanticIndex = 0;
		polygonLayout[2].Format = DXGI_FORMAT_R32G32_FLOAT;
		polygonLayout[2].InputSlot = 0;
		polygonLayout[2].AlignedByteOffset = D3D11_APPEND_ALIGNED_ELEMENT;
		polygonLayout[2].InputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;
		polygonLayout[2].InstanceDataStepRate = 0;

		polygonLayout[3].SemanticName = "NORMAL";
		polygonLayout[3].SemanticIndex = 0;
		polygonLayout[3].Format = DXGI_FORMAT_R32G32B32_FLOAT;
		polygonLayout[3].InputSlot = 0;
		polygonLayout[3].AlignedByteOffset = D3D11_APPEND_ALIGNED_ELEMENT;
		polygonLayout[3].InputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;
		polygonLayout[3].InstanceDataStepRate = 0;

		unsigned int numElements = sizeof(polygonLayout) / sizeof(polygonLayout[0]);

		// Create the vertex input layout.
		HRESULT result = nDevice->GetDevice()->CreateInputLayout(polygonLayout, 
													numElements, 
													mVertexShaderBuffer->GetBufferPointer(), 
													mVertexShaderBuffer->GetBufferSize(), 
													&mLayout);

		// Release the vertex shader buffer and pixel shader buffer since they are no longer needed.
		mVertexShaderBuffer->Release();
		mVertexShaderBuffer = NULL;

		mPixelShaderBuffer->Release();
		mPixelShaderBuffer = NULL;

		// Create a texture sampler state description.
		D3D11_SAMPLER_DESC samplerDesc;
		samplerDesc.Filter = D3D11_FILTER_MIN_MAG_MIP_LINEAR;
		samplerDesc.AddressU = D3D11_TEXTURE_ADDRESS_WRAP;
		samplerDesc.AddressV = D3D11_TEXTURE_ADDRESS_WRAP;
		samplerDesc.AddressW = D3D11_TEXTURE_ADDRESS_WRAP;
		samplerDesc.MipLODBias = 0.0f;
		samplerDesc.MaxAnisotropy = 1;
		samplerDesc.ComparisonFunc = D3D11_COMPARISON_ALWAYS;
		samplerDesc.BorderColor[0] = 0;
		samplerDesc.BorderColor[1] = 0;
		samplerDesc.BorderColor[2] = 0;
		samplerDesc.BorderColor[3] = 0;
		samplerDesc.MinLOD = 0;
		samplerDesc.MaxLOD = D3D11_FLOAT32_MAX;

		// Create the texture sampler state.
		result = nDevice->GetDevice()->CreateSamplerState(&samplerDesc, &mSampleState);
		CHECK(result);

		// Setup the description of the light dynamic constant buffer that is in the pixel shader.
		// Note that ByteWidth always needs to be a multiple of 16 if using D3D11_BIND_CONSTANT_BUFFER or CreateBuffer will fail.
		D3D11_BUFFER_DESC lightBufferDesc;
		lightBufferDesc.Usage = D3D11_USAGE_DYNAMIC;
		lightBufferDesc.ByteWidth = sizeof(LightBufferType);
		lightBufferDesc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
		lightBufferDesc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
		lightBufferDesc.MiscFlags = 0;
		lightBufferDesc.StructureByteStride = 0;

		mLightBuffer = MY_NEW(Memory::HID_Rendering, "LightShader::LightBuffer") Buffer();
		mLightBuffer->Initialize(lightBufferDesc);

		// Setup the description of the camera dynamic constant buffer that is in the vertex shader.
		D3D11_BUFFER_DESC cameraBufferDesc;
		cameraBufferDesc.Usage = D3D11_USAGE_DYNAMIC;
		cameraBufferDesc.ByteWidth = sizeof(CameraBufferType);
		cameraBufferDesc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
		cameraBufferDesc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
		cameraBufferDesc.MiscFlags = 0;
		cameraBufferDesc.StructureByteStride = 0;
	
		mCameraBuffer = MY_NEW(Memory::HID_Rendering, "LightShader::CameraBuffer") Buffer();
		mCameraBuffer->Initialize(cameraBufferDesc);

		return true;
	}
	return false;
}

void LightShader::DeinitializeShader()
{
	// Release the matrix constant buffer.
	if(mMatrixBuffer)
	{
		mMatrixBuffer->Release();
		mMatrixBuffer = NULL;
	}

	// Release the layout.
	if(mLayout)
	{
		mLayout->Release();
		mLayout = NULL;
	}

	// Release the pixel shader.
	if(mPixelShader)
	{
		mPixelShader->Release();
		mPixelShader = NULL;
	}

	// Release the vertex shader.
	if(mVertexShader)
	{
		mVertexShader->Release();
		mVertexShader = NULL;
	}

	if(mLightBuffer)
	{
		MY_DELETE(mLightBuffer);
		mLightBuffer = NULL;
	}

	if(mCameraBuffer)
	{
		MY_DELETE(mCameraBuffer);
		mCameraBuffer = NULL;
	}
}

bool LightShader::SetShaderParameters(D3DXMATRIX worldMatrix, D3DXMATRIX viewMatrix, D3DXMATRIX projectionMatrix, Camera* camera, Light* light, Texture* texture)
{
	HRESULT result;
	D3D11_MAPPED_SUBRESOURCE mappedResource;
	unsigned int bufferNumber;
	MatrixBufferType* dataPtr;
	LightBufferType* dataPtr2;
	CameraBufferType* dataPtr3;

	// Transpose the matrices to prepare them for the shader.
	D3DXMatrixTranspose(&worldMatrix, &worldMatrix);
	D3DXMatrixTranspose(&viewMatrix, &viewMatrix);
	D3DXMatrixTranspose(&projectionMatrix, &projectionMatrix);

	Graphics::D3DDevice* device = GetPlatform()->GetDevice()->As<Graphics::D3DDevice>();

	// Lock the constant buffer so it can be written to.
	result = device->GetDeviceContext()->Map(mMatrixBuffer, 0, D3D11_MAP_WRITE_DISCARD, 0, &mappedResource);
	if(FAILED(result))
	{
		return false;
	}

	// Get a pointer to the data in the constant buffer.
	dataPtr = (MatrixBufferType*)mappedResource.pData;

	// Copy the matrices into the constant buffer.
	dataPtr->World = worldMatrix;
	dataPtr->View = viewMatrix;
	dataPtr->Projection = projectionMatrix;

	// Unlock the constant buffer.
	device->GetDeviceContext()->Unmap(mMatrixBuffer, 0);

	// Set the position of the constant buffer in the vertex shader.
	bufferNumber = 0;

	// Now set the constant buffer in the vertex shader with the updated values.
	device->GetDeviceContext()->VSSetConstantBuffers(bufferNumber, 1, &mMatrixBuffer);

	// Set shader texture resource in the pixel shader.
	ID3D11ShaderResourceView* d3dTexture = texture->GetResource();
	device->GetDeviceContext()->PSSetShaderResources(0, 1, &d3dTexture);

	// Lock the camera constant buffer so it can be written to.
	ID3D11Resource* resource = mCameraBuffer->GetResource();
	result = device->GetDeviceContext()->Map(resource, 0, D3D11_MAP_WRITE_DISCARD, 0, &mappedResource);
	CHECK(result);

	// Get a pointer to the data in the constant buffer.
	dataPtr3 = (CameraBufferType*)mappedResource.pData;

	//// Copy the camera position into the constant buffer.
	dataPtr3->CameraPosition = Utility::NeroToDXD(camera->GetPosition());
	dataPtr3->Padding = 0.0f;

	//// Unlock the camera constant buffer.
	device->GetDeviceContext()->Unmap(resource, 0);

	// Set the position of the camera constant buffer in the vertex shader.
	bufferNumber = 1;

	// Now set the camera constant buffer in the vertex shader with the updated values.
	ID3D11Buffer* cameraBuffer = mCameraBuffer->GetResource();
	device->GetDeviceContext()->VSSetConstantBuffers(bufferNumber, 1, &cameraBuffer);

	// Lock the light constant buffer so it can be written to.
	result = device->GetDeviceContext()->Map(mLightBuffer->GetResource(), 0, D3D11_MAP_WRITE_DISCARD, 0, &mappedResource);
	CHECK(result);

	// Get a pointer to the data in the constant buffer.
	dataPtr2 = (LightBufferType*)mappedResource.pData;

	// Copy the lighting variables into the constant buffer.
	dataPtr2->AmbientColor		= light->GetLightBufferInfo().AmbientColor;
	dataPtr2->DiffuseColor		= light->GetLightBufferInfo().DiffuseColor;
	dataPtr2->Direction			= light->GetLightBufferInfo().Direction;
	dataPtr2->SpecularColor		= light->GetLightBufferInfo().SpecularColor;
	dataPtr2->SpecularPower		= light->GetLightBufferInfo().SpecularPower;

	// Unlock the constant buffer.
	device->GetDeviceContext()->Unmap(mLightBuffer->GetResource(), 0);

	// Set the position of the light constant buffer in the pixel shader.
	bufferNumber = 0;

	// Finally set the light constant buffer in the pixel shader with the updated values.
	ID3D11Buffer* lightBuffer = mLightBuffer->GetResource();
	device->GetDeviceContext()->PSSetConstantBuffers(bufferNumber, 1, &lightBuffer);

	return true;
}